import numpy as np
import cv2
from flask import Flask, request, jsonify, send_from_directory
import logging
import json

from services import ModelService, YOLOModelWrapper, YOLOPreprocessor, YOLOPostprocessorWithClass, MutilateModelWrapper

# 设置日志记录
logging.basicConfig(level=logging.INFO)
app = Flask(__name__)

# 初始化模型服务
model_service = ModelService()
# 四种缺陷一起检测
model_service.register_model("lixin-front", YOLOModelWrapper(model_path="runs/detect/train4/weights/best.onnx", preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024), postprocessor=YOLOPostprocessorWithClass()))

# 四种缺陷分开单独检测(单独的模型)
mutilate_model_wrapper = MutilateModelWrapper(model_configs={}, preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024))
# 起泡
mutilate_model_wrapper.addModels("bubble", YOLOModelWrapper(model_path="runs/detect/train1/weights/best.onnx", preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024)), ["缺角"])
# 缺角
mutilate_model_wrapper.addModels("chipping", YOLOModelWrapper(model_path="runs/detect/train2/weights/best.onnx", preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024)), ["缺角"])
# 裂缝
mutilate_model_wrapper.addModels("crack", YOLOModelWrapper(model_path="runs/detect/train3/weights/best.onnx", preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024)), ["缺角"])
# 砂眼
mutilate_model_wrapper.addModels("pit", YOLOModelWrapper(model_path="runs/detect/train4/weights/best.onnx", preprocessor=YOLOPreprocessor(target_width=1280, target_height=1024)), ["缺角"])
model_service.register_model("lixin-front-new", mutilate_model_wrapper)

@app.route("/api/predict", methods=["POST"])
def predict():
    try:
        if "file" not in request.files:
            return jsonify({"error": "No file provided"}), 400
        
        model_id = request.form.get("model", "default")
        detect_options = []
        detect_options_str = request.form.get('detectOptions')
    
        if not detect_options_str:
            detect_options = []
        else:
            try:
                # 尝试解析为 JSON 数组（例如: ["chipping", "crack"]）
                detect_options = json.loads(detect_options_str)
                if not isinstance(detect_options, list):
                    raise ValueError("不是有效的列表格式")
            except json.JSONDecodeError:
                # 若解析失败，尝试用逗号分割字符串（例如: "chipping,crack"）
                detect_options = [opt.strip() for opt in detect_options_str.split(',')]
        
        # 处理检测选项
        print(f"检测选项: {detect_options}")
        
        file = request.files["file"]
        if file.filename == '':
            return jsonify({"error": "Empty filename"}), 400
        
        image = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
        if image is None:
            return jsonify({"error": "Failed to decode image"}), 400
        
        options = {
            "detectOptions": ["chipping", "crack"]
        }
        
        # 使用默认模型进行预测
        result = model_service.predict(model_id, image, options)
        return jsonify([result])
        
    except Exception as e:
        logging.error(f"Error occurred during prediction: {e}")
        return jsonify({"error": str(e)}), 500

@app.route("/")
def index():
    return send_from_directory("./webapp/", "index.html")

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000)
